/*
 * This code implements stVSSIM VQA as suggested by Anush Moorthy and A.C Bovik
 * in "Efficient Motion Weighted Spatio-Temporal Video SSIM Index"
 * http://live.ece.utexas.edu/publications/2010/moorthy_spie_jan10.pdf
 *
 * Please cite the following paper in any published work if you use this software.
 * Anus K. Moorthy, Alan C. Bovik,"Efficient Motion Weighted Spatio-Temporal
 * Video SSIM Index", SPIE Conference on Human Vision and Electronic Imaging,
 * San Jose, CA, January 17-21, 2010
 *
 * File: vssim.cpp
 *
 * Author: Aditee Shrotre
 * Usage: ./vssim <ref_video.yuv> <distorted_video.yuv> <width> <height> <format>
 *      		- both the video resolutions should be equal
 *      		- The videos should be longer than 17 frames
 *      		- <format>: supported Video formats planar Y_ONLY,YUV420,YUV422,YUV444
 *	   OR ./vssim -h
 *				- for help
 *
 */

#include "vssim.h"

vssim_global_t global;
vssim_video_t refVideo;
vssim_video_t distortedVideo;

extern vssim_result_t getSSIM(int cntframe);
extern vssim_result_t getMSSIM(int cntframe);

/*
 * This function finds the max value in the given image structure
 */
uint32_t findMax(IplImage *img)
{
	unsigned char max;
	uint64_t i, size = img->width*img->height;

	max = (unsigned char)img->imageData[0];

	for (i = 0; i < size; i++)
	{
		if ((unsigned char)img->imageData[i] > max)
			max = (unsigned char)img->imageData[i];
	}

	return (max);
}

/*
 * This function calculates the percentile SSIM.
 * The implementation of the P-SSIM has been optimized for performance. Instead
 * of sorting the values in the entire image the function accumulates the lowest
 * 6% scores in the image and then averages them over cardinality
 */
float calcPSSIM(IplImage *ssim_map)
{
	IplImage *ssim_clone;
	ssim_clone = cvCloneImage(ssim_map);

	CvScalar avg_ssim = cvAvg(ssim_clone);

	/* Cardinality of the SSIM values to be considered */
	uint32_t cardinality = (int)((float)0.06 * (float)(ssim_map->width*ssim_map->height));

	uint64_t size = ssim_clone->width*ssim_clone->height;

	unsigned char i;
	uint32_t cnt = 0;
	uint64_t sum = 0;
	float avg = 0.0;

	/* start with the average value, but if avg value is 0, then consider 1
	 * as the starting point */
	float value = ((uint32_t)avg_ssim.val[0] <= 0) ? (1) : ((uint32_t)avg_ssim.val[0]);
	for(i = 0; i < value; i++)
	{
		uint64_t j;
		cnt = 0;
		sum = 0;
		for (j = 0; j < size; j++)
		{
			/* Accumulate the values that are less than the average of the image */
			if ((unsigned char)ssim_clone->imageData[j] <= i)
			{
				sum += (unsigned char)ssim_clone->imageData[j];
				cnt++;
				if (cnt == cardinality)
					break;
			}
		}
		/* If the cnt reaches cardinality average for P-SSIM score */
		if (cnt == cardinality)
		{
			avg = (float)sum/(float)cnt;
			break;
		}
		/* If the cardinality is not met then increment the upper threshold of
		 * the value */
		else if (cnt < cardinality)
			value += 10;	// increment the upper limit to get to the cardinality
	}

	return avg;
}

/*
 * Prints the usage for the code
 */
static void print_usage()
{
	printf("USAGE: \n");
	printf("\t./vssim <ref_video.yuv> <distorted_video.yuv> <width> <height> <format>\n");
	printf("\t\t- <ref_video.yuv>. reference video file in supported format\n");
	printf("\t\t- <distorted_video.yuv>. Distorted video file in supported format\n");
	printf("\t\t- <width> <height> width and height of the video files. (Both the videos "
			"are expected to have the same resolution\n");
	printf("\t\t- <format>: one of the supported planar formats: Y_ONLY/YUV420/YUV422/YUV444 "
			"(Both the videos are expected to have the same format)");

	return;
}

/*
 * This function parses the input arguments
 */
static void parse_args(int argc, char *argv[])
{
	/* get the reference video file name */
	strcpy(refVideo.fileName, argv[1]);

	/* get the distorted video file name */
	strcpy(distortedVideo.fileName, argv[2]);

	/* get the resolution of the video frames */
	global.size.width = atoi(argv[3]);
	global.size.height = atoi(argv[4]);

	/* get the format of the video frames */
	if ((strcmp(argv[5], "Y_ONLY")) == 0)
		global.format = FORMAT_YONLY;
	else if ((strcmp(argv[5], "YUV420")) == 0)
		global.format = FORMAT_YUV420;
	else if ((strcmp(argv[5], "YUV422")) == 0)
		global.format = FORMAT_YUV422;
	else if ((strcmp(argv[5], "YUV444")) == 0)
		global.format = FORMAT_YUV444;
	else
		printf("Format not supported\n");

	return;
}

/*
 * allocates the global memory
 */
static void allocmem()
{
	global.Sframe = (vssim_FrScore_t *)malloc(sizeof(vssim_FrScore_t) * (global.nframes/16));
	global.Tframe = (vssim_FrScore_t *)malloc(sizeof(vssim_FrScore_t) * (global.nframes/16));

	return;
}

/*
 * Frees all the allocated memory
 */
static void freemem()
{
	free(global.Sframe);
	free(global.Tframe);

	return;
}

/*
 * This function averages the Tframe and the Sframe scores for Y,U and V seperately
 * and then combines them with weights, where summation of the weights = 1
 */
static float avgScore(vssim_FrScore_t *scores, int nframes)
{
	float Yavg = 0.0;
	float Uavg = 0.0;
	float Vavg = 0.0;
	float score = 1.0;

	for (int i = 0; i < nframes; i++)
	{
		/* Accumulate the scores */
		Yavg += scores[i].score_Y;
		Uavg += scores[i].score_U;
		Vavg += scores[i].score_V;
	}
	/* Average over the frames for which the indexes are calculated */
	Yavg /= nframes;
	Uavg /= nframes;
	Vavg /= nframes;

	/* Combine the Y,U,V scores */
	score = 0.8 * Yavg + 0.1 * Uavg + 0.1 * Vavg;

	return score;
}

/*
 * This is the main entry point to the application
 * main parses the arguments and calls the functions to calculate the SS-SSIM
 * and temporal SSIM. At the end it computes the final scores and prints.
 */
int main (int argc, char *argv[])
{
	vssim_result_t result;
	long size;
	int cntframe = 0;

	/* Parse the arguments */
	if (argc < 6 || (!strcmp(argv[1], "-h")))
	{
		print_usage();
		return FAIL;
	}
	else
	{
		parse_args(argc, argv);
	}

	/* Open the reference video file */
	refVideo.fptr = fopen(refVideo.fileName, "r");
	if (refVideo.fptr == NULL)
	{
		printf("Could not open the Reference Video file. Exiting..\n");
		return FAIL;
	}

	/* The videos are expected to have more than 17 frames */
	fseek(refVideo.fptr, 0, SEEK_END);
	size = ftell(refVideo.fptr);
	fseek(refVideo.fptr, 0, SEEK_SET);
	if (size < (17 * FRAME_HEIGHT(&global) * global.size.width))
	{
		printf("Reference Video does not have enough frames. # of frames in the "
				"video are expected to be greater than 17\n");
		return FAIL;
	}

	/* open the distorted video file */
	distortedVideo.fptr = fopen(distortedVideo.fileName, "r");
	if (distortedVideo.fptr == NULL)
	{
		printf("Could not open the Distorted Video file. Exiting..\n");
		return FAIL;
	}

	fseek(distortedVideo.fptr, 0, SEEK_END);
	size = ftell(distortedVideo.fptr);
	fseek(distortedVideo.fptr, 0, SEEK_SET);
	if (size < (17 * FRAME_HEIGHT(&global) * global.size.width))
	{
		printf("Distorted Video does not have enough frames. # of frames in "
				"the video are expected to be greater than 17\n");
		return -1;
	}

	printf("height y: %d, height u: %d, height v: %d, frame height: %d\n",
			HEIGHT_Y(&global), HEIGHT_U(&global), HEIGHT_V(&global), FRAME_HEIGHT(&global));

	/* total number of frames = size of file / size of frames */
	int nrefFrames, ndistFrames;
	nrefFrames = (size / (FRAME_HEIGHT(&global)*global.size.width));
	ndistFrames = (size / (FRAME_HEIGHT(&global)*global.size.width));
	global.nframes = (nrefFrames < ndistFrames) ? nrefFrames : ndistFrames;

	printf("Total Number of Frames: %d\n", global.nframes);

	allocmem();

	cntframe+=16;
	do
	{
		/* Calculate the SS-SSIM */
		printf("Calculating SS-SSIM...\n");
		result = getSSIM(cntframe);

		/* Calculate the temporal SSIM */
		printf("Calculating MC-VSSIM...\n");
		result = getMSSIM(cntframe);

		/* get every 16th frame */
		cntframe+=16;

	} while(cntframe <= global.nframes);

	/* The mean of the S and T scores */
	float Sfinal, Tfinal, stSSIM;

	Sfinal = avgScore(global.Sframe, (global.nframes/16));
	printf("Avg S-Score: %f\n", Sfinal);

	Tfinal = avgScore(global.Tframe, (global.nframes/16));
	printf("Avg T-Score: %f\n", Tfinal);

	/* Final Score = Tvideo * Svideo */
	stSSIM = Sfinal * Tfinal;
	printf("stVSSIM Score: %f\n", stSSIM);

	fclose(refVideo.fptr);
	fclose(distortedVideo.fptr);

	/* Clean up the memory */
	freemem();

	return result;
}
